import sqlite3
import os
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D

def get_db_files(path):
    """获取指定路径下所有的db文件"""
    db_files = [f for f in os.listdir(path) if f.endswith('.db')]
    return db_files

def read_data_from_db(db_file):
    """从db文件中读取所有表的数据"""
    conn = sqlite3.connect(db_file)
    cursor = conn.cursor()

    tables = ["sample_time", "decode_attention", "decode_gemm", "decode_moe", 
              "prefill_attention", "prefill_gemm", "prefill_moe"]

    data = {}
    for table in tables:
        cursor.execute(f"SELECT * FROM {table}")
        data[table] = cursor.fetchall()
    
    conn.close()
    return data

def plot_scatter(path,x, y, title, xlabel, ylabel):
    """绘制散点图"""
    plt.figure()
    plt.scatter(x, y)
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.savefig(os.path.join(path,f'{title}.png'))
    plt.close()

def plot_surface(path, x, y, z, title, xlabel, ylabel, zlabel):
    X, Y = np.meshgrid(np.unique(x), np.unique(y))  # 创建网格
    Z = np.full(X.shape, np.nan)  # 创建一个空的 Z 数组

    # 使用字典来合并重复的 (x, y) 对
    data_dict = {}
    for i in range(len(z)):
        x_idx = np.where(np.unique(x) == x[i])[0][0]
        y_idx = np.where(np.unique(y) == y[i])[0][0]

        # 合并重复值（计算平均值）
        key = (x[i], y[i])
        if key not in data_dict:
            data_dict[key] = []
        data_dict[key].append(z[i])

    # 填充 Z 数组，合并重复的 (x, y)
    for (x_val, y_val), values in data_dict.items():
        x_idx = np.where(np.unique(x) == x_val)[0][0]
        y_idx = np.where(np.unique(y) == y_val)[0][0]
        Z[y_idx, x_idx] = np.mean(values)  # 取平均值填充

    # 跳过缺失数据，直接丢弃包含 NaN 的 (x, y) 坐标
    missing_indices = np.argwhere(np.isnan(Z))
    if len(missing_indices) > 0:
        print(f"Skipping the following missing (x, y) coordinates: {missing_indices}")

    # 绘制曲面图
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    surf = ax.plot_surface(X, Y, Z, cmap='viridis')  # 绘制曲面图
    ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_zlabel(zlabel)

    # 设置 color bar
    fig.colorbar(surf, ax=ax, shrink=0.5, aspect=5)

    # 保存图像
    plt.savefig(os.path.join(path, f'{title}.png'))
    plt.close()

def process_data_for_table(path,table_name, data, tp_size, pp_size):
    """根据表名处理数据并绘图"""
    if table_name == "sample_time":
        # 只有一个因变量，绘制散点图
        batchsize = [row[2] for row in data]
        sample_time = [row[1]/1e6 for row in data]
        plot_scatter(path,batchsize, sample_time, f"{table_name}_tp{tp_size}_pp{pp_size}",
                     "Batch Size", "Sample Time")
    
    elif table_name == "decode_attention":
        # 取 current_length = 10 的数据，绘制曲面图
        filtered_data = [row for row in data if row[3] == 10]
        if filtered_data:
            batchsize = [row[2] for row in filtered_data]
            input_length = [row[4] for row in filtered_data]
            decode_attention = [row[1]/1e6 for row in filtered_data]

            # 绘制曲面图
            plot_surface(path,input_length, batchsize, decode_attention,
                         f"{table_name}_tp{tp_size}_pp{pp_size}", 
                         "Input Length", "Batch Size", "Decode Attention")
        filtered_data_batchsize=[row for row in data if row[3]==10 and row[4]==256]
        if filtered_data_batchsize:
            batchsize = [row[2] for row in filtered_data]
            decode_attention = [row[1]/1e6 for row in filtered_data]
            plot_scatter(path,batchsize,decode_attention , f"{table_name}_tp{tp_size}_pp{pp_size}_input=256",
                     "Batch Size", "Sample Time")
    elif table_name == "decode_gemm":
        batchsize = [row[2] for row in data]
        decode_gemm = [row[1]/1e6 for row in data]
        plot_scatter(path,batchsize, decode_gemm, f"{table_name}_tp{tp_size}_pp{pp_size}",
                     "Batch Size", "Decode GEMM")
    
    elif table_name == "decode_moe":
        batchsize = [row[2] for row in data]
        decode_moe = [row[1]/1e6 for row in data]
        plot_scatter(path,batchsize, decode_moe, f"{table_name}_tp{tp_size}_pp{pp_size}",
                     "Batch Size", "Decode MOE")
    
    elif table_name == "prefill_attention":
        input_tokens = [row[2] for row in data]
        prefill_attention = [row[1]/1e6 for row in data]
        plot_scatter(path,input_tokens, prefill_attention, f"{table_name}_tp{tp_size}_pp{pp_size}",
                     "Input Tokens", "Prefill Attention")
    
    elif table_name == "prefill_gemm":
        input_tokens = [row[2] for row in data]
        prefill_gemm = [row[1]/1e6 for row in data]
        plot_scatter(path,input_tokens, prefill_gemm, f"{table_name}_tp{tp_size}_pp{pp_size}",
                     "Input Tokens", "Prefill GEMM")
    
    elif table_name == "prefill_moe":
        input_tokens = [row[2] for row in data]
        prefill_moe = [row[1]/1e6 for row in data]
        plot_scatter(path,input_tokens, prefill_moe, f"{table_name}_tp{tp_size}_pp{pp_size}",
                     "Input Tokens", "Prefill MOE")

def main(path):
    db_files = get_db_files(path)
    
    for db_file in db_files:
        # 提取tp_size和pp_size
        parts = db_file.split('_')
        tp_size = parts[1][2:]
        pp_size = parts[2][2:]
        
        data = read_data_from_db(os.path.join(path, db_file))
        
        for table_name, table_data in data.items():
            if table_data:  # 如果表有数据
                process_data_for_table(path,table_name, table_data, tp_size, pp_size)

if __name__ == "__main__":
    # 指定DB文件所在的路径
    db_path = "/root/llm/framework/MixFrame/mixserve/benchdb"
    main(db_path)
